# Nikhil Ketkar, Wed Feb 20 03:34:15 PST 2008


class Particle(object):
    def __init__(self,symbol,charge,number):
        self.symbol = symbol
        self.charge = charge
        self.number = number

    def __str__(self):
        return self.symbol

    def __mul__(self, operand):
        return tuple([self for i in xrange(0, operand)])
                      
    def __rmul__(self, operand):
        return tuple([self for i in xrange(0, operand)])
            
    def __add__(self, operand):
        if hasattr(operand,'__iter__'):
            temp = list(operand)
            temp.append(self)
            return tuple(temp)
        else:
            return (self,operand)
        
    def __radd__(self,operand):
        if hasattr(operand,'__iter__'):
            temp = list(operand)
            temp.append(self)
            return tuple(temp)
        else:
            return (self,operand)

class Nucleus(Particle):
    def __str__(self):
        return "(" + str(self.number) + ")" + str(self.symbol)

class UnbalancedCharge:
    pass

class UnbalancedNumber:
    pass

def general_append(collection, seperator):
    string = ""
    for item in collection:
        string += item.__str__()
        string += seperator
    return string[:len(string) - len(seperator)]
    
def general_add(collection, attribute):
    sum = 0
    for item in collection:
        sum += item.__dict__[attribute]
    return sum

def remove_duplicates(given_list1, given_list2):
    for index1, item1 in enumerate(given_list1):
        for index2, item2 in enumerate(given_list2):
            if not item1 is None and not item2 is None and item1 == item2:
                given_list1[index1] = given_list2[index2] = None
                break
    return ([item for item in given_list1 if not item is None], 
            [item for item in given_list2 if not item is None])

def group_list(given_list):
    return_dict = {}
    for item in given_list:
        if return_dict.has_key(item):
            return_dict[item] += 1
        else:
            return_dict[item] = 1
    return_string = ""
    for index, item in return_dict.iteritems():
        if item > 1:
            return_string += str(item)
            return_string += " "
        return_string += index.__str__()
        return_string += " + "
    return return_string[:len(return_string) - 3]

class Reaction(object):
    def __init__(self,reactants, products, chain_id = None):
        if general_add(reactants, "charge") != general_add(products, "charge"):
            raise UnbalancedCharge
        elif general_add(reactants, "number") != general_add(products, "number"):
            raise UnbalancedNumber
        else:
            self.reactants = reactants
            self.products = products

        if isinstance(chain_id, ChainReaction):
            chain_id.add(self)

    def __str__(self):
        return general_append(self.reactants, " + ") \
            + " -> " \
            + general_append(self.products," + ")

class ChainReaction:
    def __init__(self, symbol):
        self.symbol = symbol
        self.reactions = []
    def add(self, reaction):
        self.reactions.append(reaction)
    def __str__(self):
        reactants = []
        products = []
        for reaction in self.reactions:
            for reactant in reaction.reactants:
                reactants.append(reactant) 
            for product in reaction.products:
                products.append(product) 
        net_reactants, net_products = remove_duplicates(reactants, products)
        grouped_reactants = group_list(net_reactants)
        grouped_products = group_list(net_products)        
        return self.symbol + "chain: \n" \
            + general_append(self.reactions, "\n") + "\n" \
            + grouped_reactants + " -> " +  grouped_products 
            
if __name__ == "__main__":
    
    print """Executing...
em = Particle("e-", -1, 0)
ep = Particle("e+", 1, 0)
p = Particle("p", 1, 1)
n = Particle("n", 0, 1)
neutrino = Particle("nu_e", 0, 0)
gamma = Particle("gamma", 0, 0) """

    em = Particle("e-", -1, 0)
    ep = Particle("e+", 1, 0)
    p = Particle("p", 1, 1)
    n = Particle("n", 0, 1)
    neutrino = Particle("nu_e", 0, 0)
    gamma = Particle("gamma", 0, 0)
    
    print ""
    print "Output of \"print em, ep, p, n, neutrino, gamma\""
    print em, ep, p, n, neutrino, gamma
    print ""

    d = Nucleus("H", 1, 2)
    li6 = Nucleus("Li", 3, 6)
    he4 = Nucleus("He", 2, 4)
 
    print "Output of \"print d, li6, he4\""
    print d, li6, he4
    print ""

    print "Output of \"print Reaction((li6, d), (he4, he4))\""
    print Reaction((li6, d), (he4, he4))
    print ""

    print "Checking if exceptions are raised for charge and number imbalance"
    try:
        Reaction((li6, d), (he4, he4, he4))
        print "Failed charge test."
    except UnbalancedCharge:
        print "Success charge test."

    try:
        Reaction((li6, d), (he4, he4, n))
        print "Failed number test."
    except UnbalancedNumber:
        print "Success number test."
    print ""
    
    print "Output of \"print Reaction((li6 + d), (he4 + he4))\""    
    print Reaction((li6 + d), (he4 + he4))
    print ""

    print "Output of \"Reaction((li6 + d + d), (he4 + he4 + d))\""
    print Reaction((li6 + d + d), (he4 + he4 + d))
    print ""

    print "Output of \"Reaction((li6 + d + d + d + d), (he4 + he4 + d + d + d))\""
    print Reaction((li6 + d + d + d + d), (he4 + he4 + d + d + d))
    print ""
    
    print "Output of \"Reaction((2 * d * 4), (2 * d * 2 * 2))\""
    print Reaction((2 * d * 4), (2 * d * 2 * 2))
    print ""

    print """Executing...
he3 = Nucleus("He", 2, 3)
chnPP = ChainReaction("proton-proton (branch I)")
Reaction((p,p), (d, ep, neutrino), chnPP)
Reaction((p,p), (d, ep, neutrino), chnPP)
Reaction((d,p), (he3, gamma), chnPP)
Reaction((d,p), (he3, gamma), chnPP)
Reaction((he3,he3), (he4, p, p), chnPP)"""

    print ""
    he3 = Nucleus("He", 2, 3)
    chnPP = ChainReaction("proton-proton (branch I)")
    Reaction((p,p), (d, ep, neutrino), chnPP)
    Reaction((p,p), (d, ep, neutrino), chnPP)
    Reaction((d,p), (he3, gamma), chnPP)
    Reaction((d,p), (he3, gamma), chnPP)
    Reaction((he3,he3), (he4, p, p), chnPP)

    print "Output of print \"chnPP\""
    print chnPP
    print ""
